#! /usr/bin/env python3
#coding=utf-8

#460. LFU Cache

from collections import OrderedDict


class LFUCache(object):
    class FreqNode:
        def __init__(self, freq=0):
            self.map = OrderedDict()
            self.freq = freq
            self.next = None
            self.prev = None

        def connect(self, node):
            if node is None:
                self.next = None
                return
            self.next = node
            node.prev = self

    def __init__(self, capacity):
        """
        :type capacity: int
        """
        self.capacity = capacity
        self.kv_map = {}
        self.kf_map = {}
        self.freq_node_list_head = None

    def __get_cur_capacity(self):
        return self.kv_map.__len__()

    def get(self, key):
        value = self.kv_map.get(key)
        if value is None:
            return -1
        self.inc_freq(key)
        return value

    def inc_freq(self, key):
        freq_node = self.kf_map[key]
        freq_node.map.pop(key)
        next_freq_node = freq_node.next
        if next_freq_node is not None and next_freq_node.freq == freq_node.freq + 1:
            next_freq_node.map[key] = key
            self.kf_map[key] = next_freq_node
        else:
            new = self.FreqNode(freq_node.freq + 1)
            new.map[key] = key
            freq_node.connect(new)
            new.connect(next_freq_node)
            self.kf_map[key] = new
        if len(freq_node.map) == 0:
            self.__remove_freq_node(freq_node)

    def set(self, key, value):
        if self.capacity == 0:
            return
        if key in self.kv_map:
            self.kv_map[key] = value
        else:
            if self.kv_map.__len__() == self.capacity:
                self.remove_old()
            self.kv_map[key] = value
            freq_node = self.FreqNode(0)
            freq_node.map[key] = key
            self.kf_map[key] = freq_node
            self.add_to_head(freq_node)
        self.inc_freq(key)

    def add_to_head(self, node):
        if self.freq_node_list_head is None:
            self.freq_node_list_head = node
        else:
            node.connect(self.freq_node_list_head)
            self.freq_node_list_head = node

    def __remove_freq_node(self, node):
        if self.freq_node_list_head is node:
            self.freq_node_list_head = node.next
            self.freq_node_list_head.prev = None
        else:
            node.prev.connect(node.next)

    def remove_old(self):
        cap = self.__get_cur_capacity()
        assert(cap > 0)
        head = self.freq_node_list_head
        key = head.map.popitem(last=False)[0]
        self.kf_map.pop(key)
        self.kv_map.pop(key)
        if head.map.__len__() == 0:
            self.freq_node_list_head = head.next

cache = LFUCache(2)

cache.set(1, 1)
cache.set(2, 2)
print(cache.get(1))
cache.set(3, 3)
print(cache.get(2))
print(cache.get(3))
cache.set(4, 4)
print(cache.get(1))
print(cache.get(3))
print(cache.get(4))


